// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <optional>
#include <pollux/testing/functions/prestosql/function_base_test.h>
#include <pollux/functions/prestosql/types/timestamp_with_time_zone_type.h>

using namespace kumo::pollux;
using namespace kumo::pollux::test;
using namespace kumo::pollux::functions::test;

namespace {

class ArrayMaxTest : public FunctionBaseTest {
 protected:
  void testArrayMax(const VectorPtr& input, const VectorPtr& expected) {
    auto result = evaluate<BaseVector>("array_max(C0)", make_row_vector({input}));
    assertEqualVectors(expected, result);
  }
};

TEST_F(ArrayMaxTest, booleanWithNulls) {
  auto input = make_nullable_array_vector<bool>(
      {{true, false},
       {true},
       {false},
       {},
       {true, false, true, std::nullopt},
       {std::nullopt, true, false, true},
       {false, false, false},
       {true, true, true}});

  auto expected = make_nullable_flat_vector<bool>(
      {true,
       true,
       false,
       std::nullopt,
       std::nullopt,
       std::nullopt,
       false,
       true});

  testArrayMax(input, expected);
}

TEST_F(ArrayMaxTest, booleanNoNulls) {
  auto input = make_nullable_array_vector<bool>(
      {{true, false},
       {true},
       {false},
       {},
       {false, false, false},
       {true, true, true}});

  auto expected = make_nullable_flat_vector<bool>(
      {true, true, false, std::nullopt, false, true});

  testArrayMax(input, expected);
}

TEST_F(ArrayMaxTest, varcharWithNulls) {
  auto input = make_nullable_array_vector<StringView>({
      {"red"_sv, "blue"_sv},
      {std::nullopt, "blue"_sv, "yellow"_sv, "orange"_sv},
      {},
      {"red"_sv, "purple"_sv, "green"_sv},
  });

  auto expected = make_nullable_flat_vector<StringView>(
      {"red"_sv, std::nullopt, std::nullopt, "red"_sv});

  testArrayMax(input, expected);
}

TEST_F(ArrayMaxTest, varcharNoNulls) {
  auto input = make_array_vector<StringView>({
      {"red"_sv, "blue"_sv},
      {"blue"_sv, "yellow"_sv, "orange"_sv},
      {},
      {"red"_sv, "purple"_sv, "green"_sv},
  });

  auto expected = make_nullable_flat_vector<StringView>(
      {"red"_sv, "yellow"_sv, std::nullopt, "red"_sv});

  testArrayMax(input, expected);
}

// Test non-inlined (> 12 length) nullable strings.
TEST_F(ArrayMaxTest, longVarcharWithNulls) {
  auto input = make_nullable_array_vector<StringView>({
      {"red shiny car ahead"_sv, "blue clear sky above"_sv},
      {std::nullopt,
       "blue clear sky above"_sv,
       "yellow rose flowers"_sv,
       "orange beautiful sunset"_sv},
      {},
      {"red shiny car ahead"_sv,
       "purple is an elegant color"_sv,
       "green plants make us happy"_sv},
  });

  auto expected = make_nullable_flat_vector<StringView>(
      {"red shiny car ahead"_sv,
       std::nullopt,
       std::nullopt,
       "red shiny car ahead"_sv});

  testArrayMax(input, expected);
}

// Test non-inlined (> 12 length) strings.
TEST_F(ArrayMaxTest, longVarcharNoNulls) {
  auto input = make_nullable_array_vector<StringView>({
      {"red shiny car ahead"_sv, "blue clear sky above"_sv},
      {"blue clear sky above"_sv,
       "yellow rose flowers"_sv,
       "orange beautiful sunset"_sv},
      {},
      {"red shiny car ahead"_sv,
       "purple is an elegant color"_sv,
       "green plants make us happy"_sv},
  });

  auto expected = make_nullable_flat_vector<StringView>(
      {"red shiny car ahead"_sv,
       "yellow rose flowers"_sv,
       std::nullopt,
       "red shiny car ahead"_sv});

  testArrayMax(input, expected);
}

TEST_F(ArrayMaxTest, timestampWithTimezone) {
  auto testArrayMax = [this](
                          const std::vector<std::optional<int64_t>>& inputArray,
                          std::optional<int64_t> expectedValue) {
    // Test with primitive types.
    auto input = make_row_vector({make_array_vector(
        {0}, make_nullable_flat_vector(inputArray, TIMESTAMP_WITH_TIME_ZONE()))});
    VectorPtr expected = make_nullable_flat_vector<int64_t>(
        {expectedValue}, TIMESTAMP_WITH_TIME_ZONE());

    auto result = evaluate<BaseVector>("array_max(C0)", input);
    assertEqualVectors(expected, result);

    // array_max does not support nulls inside complex types. To exclude these
    // we exclude tests where the expected result is null.  However, the result
    // can also be null when the input is empty, so we add an exception for that
    // case.
    if (expectedValue.has_value() || inputArray.empty()) {
      // Test wrapped in complex type.
      input = make_row_vector({make_array_vector(
          {0},
          make_row_vector({make_nullable_flat_vector(
              inputArray, TIMESTAMP_WITH_TIME_ZONE())}))});
      expected = make_row_vector(
          {expected}, [&](vector_size_t) { return inputArray.empty(); });

      result = evaluate<BaseVector>("array_max(C0)", input);
      assertEqualVectors(expected, result);
    }
  };

  testArrayMax(
      {pack(-1, 0), pack(0, 1), pack(1, 2), pack(2, 3), pack(3, 4), pack(4, 5)},
      pack(4, 5));
  testArrayMax(
      {pack(4, 0),
       pack(3, 1),
       pack(2, 2),
       pack(1, 3),
       pack(0, 4),
       pack(-1, 5),
       pack(-2, 6)},
      pack(4, 0));
  testArrayMax(
      {pack(-5, 3), pack(-4, 2), pack(-3, 1), pack(-2, 0), pack(-1, 4)},
      pack(-1, 4));
  testArrayMax(
      {pack(101, 4), pack(102, 0), pack(103, 1), pack(104, 2), pack(105, 3)},
      pack(105, 3));
  testArrayMax({}, std::nullopt);
  testArrayMax(
      {pack(101, 4), pack(102, 0), pack(103, 1), pack(104, 2), std::nullopt},
      std::nullopt);
  testArrayMax(
      {std::nullopt, pack(-1, 4), pack(-2, 5), pack(-3, 1), pack(-4, 0)},
      std::nullopt);
  testArrayMax({std::nullopt}, std::nullopt);
  testArrayMax({pack(1, 0), pack(1, 1), pack(1, 2)}, pack(1, 0));
}

// Test documented example.
TEST_F(ArrayMaxTest, docs) {
  {
    auto input = make_nullable_array_vector<int32_t>(
        {{1, 2, 3}, {-1, -2, -2}, {-1, -2, std::nullopt}, {}});
    auto expected =
        make_nullable_flat_vector<int32_t>({3, -1, std::nullopt, std::nullopt});
    testArrayMax(input, expected);
  }
  {
    static const float kNaN = std::numeric_limits<float>::quiet_NaN();
    static const float kInfinity = std::numeric_limits<float>::infinity();
    auto input = make_nullable_array_vector<float>(
        {{-1, kNaN, std::nullopt}, {-1, -2, -3, kNaN}, {kInfinity, kNaN}});
    auto expected = make_nullable_flat_vector<float>({std::nullopt, kNaN, kNaN});
    testArrayMax(input, expected);
  }
}

template <typename Type>
class ArrayMaxIntegralTest : public FunctionBaseTest {
 public:
  using T = typename Type::NativeType;

  void testArrayMax(const VectorPtr& input, const VectorPtr& expected) {
    auto result =
        evaluate<SimpleVector<T>>("array_max(C0)", make_row_vector({input}));
    assertEqualVectors(expected, result);
  }

  void testWithNulls() {
    auto input = make_nullable_array_vector<T>(
        {{std::numeric_limits<T>::min(),
          0,
          1,
          2,
          3,
          std::numeric_limits<T>::max()},
         {std::numeric_limits<T>::max(),
          3,
          2,
          1,
          0,
          -1,
          std::numeric_limits<T>::min()},
         {101, 102, 103, std::numeric_limits<T>::max(), std::nullopt},
         {std::nullopt, -1, -2, -3, std::numeric_limits<T>::min()},
         {},
         {std::nullopt}});

    auto expected = make_nullable_flat_vector<T>(
        {std::numeric_limits<T>::max(),
         std::numeric_limits<T>::max(),
         std::nullopt,
         std::nullopt,
         std::nullopt,
         std::nullopt});

    testArrayMax(input, expected);
  }

  void testNoNulls() {
    auto input = make_array_vector<T>(
        {{std::numeric_limits<T>::min(), 0, 1, 2, 3, 4},
         {std::numeric_limits<T>::max(), 3, 2, 1, 0, -1, -2},
         {std::numeric_limits<T>::max(),
          101,
          102,
          103,
          104,
          105,
          std::numeric_limits<T>::max()},
         {}});

    auto expected = make_nullable_flat_vector<T>(
        {4,
         std::numeric_limits<T>::max(),
         std::numeric_limits<T>::max(),
         std::nullopt});

    testArrayMax(input, expected);
  }
};

template <typename Type>
class ArrayMaxFloatingPointTest : public FunctionBaseTest {
 public:
  using T = typename Type::NativeType;

  void testArrayMax(VectorPtr input, VectorPtr expected) {
    auto result =
        evaluate<SimpleVector<T>>("array_max(C0)", make_row_vector({input}));
    assertEqualVectors(expected, result);
  }

  void testWithNulls() {
    auto input = make_nullable_array_vector<T>(
        {{0.0000, 0.00001},
         {std::nullopt, 1.1, 1.11, -2.2, -1.0, std::numeric_limits<T>::min()},
         {std::numeric_limits<T>::min(),
          -0.0001,
          -0.0002,
          -0.0003,
          std::numeric_limits<T>::max()},
         {},
         {std::nullopt},
         {std::numeric_limits<T>::min(), 1.1, 1.22222, 1.33, std::nullopt},
         {-0.00001, -0.0002, 0.0001}});

    auto expected = make_nullable_flat_vector<T>(
        {0.00001,
         std::nullopt,
         std::numeric_limits<T>::max(),
         std::nullopt,
         std::nullopt,
         std::nullopt,
         0.0001});
    testArrayMax(input, expected);
  }

  void testNoNulls() {
    auto input = make_array_vector<T>(
        {{0.0000, 0.00001},
         {std::numeric_limits<T>::max(),
          1.1,
          1.11,
          -2.2,
          -1.0,
          std::numeric_limits<T>::min()},
         {std::numeric_limits<T>::min(),
          -0.0001,
          -0.0002,
          -0.0003,
          std::numeric_limits<T>::max()},
         {},
         {1.1, 1.22222, 1.33},
         {-0.00001, -0.0002, 0.0001}});

    auto expected = make_nullable_flat_vector<T>(
        {0.00001,
         std::numeric_limits<T>::max(),
         std::numeric_limits<T>::max(),
         std::nullopt,
         1.33,
         0.0001});
    testArrayMax(input, expected);
  }

  void testExtremeValues() {
    static const T kNaN = std::numeric_limits<T>::quiet_NaN();
    static const T kInfinity = std::numeric_limits<T>::infinity();
    static const T kNegativeInfinity = -1 * std::numeric_limits<T>::infinity();
    auto input = make_nullable_array_vector<T>(
        {{-1, std::nullopt, kNaN},
         {-1, std::nullopt, 2},
         {-1, 0, 2},
         {kNegativeInfinity, kNegativeInfinity},
         {-1, 2, kInfinity},
         {kInfinity, kNaN}});
    auto expected = make_nullable_flat_vector<T>(
        {std::nullopt, std::nullopt, 2, kNegativeInfinity, kInfinity, kNaN});
    testArrayMax(input, expected);
  }
};

} // namespace

TYPED_TEST_SUITE(ArrayMaxIntegralTest, FunctionBaseTest::IntegralTypes);

TYPED_TEST(ArrayMaxIntegralTest, arrayMaxNullable) {
  this->testWithNulls();
}

TYPED_TEST(ArrayMaxIntegralTest, arrayMax) {
  this->testNoNulls();
}

TYPED_TEST_SUITE(
    ArrayMaxFloatingPointTest,
    FunctionBaseTest::FloatingPointTypes);

TYPED_TEST(ArrayMaxFloatingPointTest, arrayMaxNullable) {
  this->testWithNulls();
}

TYPED_TEST(ArrayMaxFloatingPointTest, arrayMax) {
  this->testNoNulls();
}

TYPED_TEST(ArrayMaxFloatingPointTest, arrayMaxExtreme) {
  this->testExtremeValues();
}
